#ifndef JPEG_TOOL_H
#define JPEG_TOOL_H

// The following codes are written based on jpeg_read.c & jpeg_write.c
#include <stdio.h>
#include <stdlib.h>
#include "include/jpeglib.h"
#include "include/jpegint.h"
#include "include/jerror.h"

void jpeg_read_done(struct jpeg_decompress_struct *srcinfo) {
  jpeg_finish_decompress(srcinfo);
  jpeg_destroy_decompress(srcinfo);
}

void jpeg_write_done(struct jpeg_compress_struct *dstinfo) {
  jpeg_finish_compress(dstinfo);
  jpeg_destroy_compress(dstinfo);
}

void jpeg_read_prepare(struct jpeg_decompress_struct *srcinfo, struct jpeg_error_mgr *srcjerr, FILE* infile) {
  srcinfo->err = jpeg_std_error(srcjerr);
  jpeg_create_decompress(srcinfo);
  jpeg_stdio_src(srcinfo, infile);
  jpeg_read_header(srcinfo, TRUE);
  switch (srcinfo->out_color_space) {
    case JCS_GRAYSCALE:
      srcinfo->out_color_components = 1;
      break;
    case JCS_RGB:
      srcinfo->out_color_components = 3;
      break;
    case JCS_YCbCr:
      srcinfo->out_color_components = 3;
      break;
    case JCS_CMYK:
      srcinfo->out_color_components = 4;
      break;
    case JCS_YCCK:
      srcinfo->out_color_components = 4;
      break;
  }
}

void jpeg_write_prepare(struct jpeg_compress_struct *dstinfo, struct jpeg_error_mgr *dstjerr, FILE* outfile) {
  dstinfo->err = jpeg_std_error(dstjerr);
  jpeg_create_compress(dstinfo);
  jpeg_stdio_dest(dstinfo, outfile);
}

jvirt_barray_ptr* jpeg_copy_data(struct jpeg_decompress_struct *srcinfo, struct jpeg_compress_struct *dstinfo) {
  int ci,n,i,j;
  jvirt_barray_ptr *srccoef_arrays = NULL;
  jvirt_barray_ptr *dstcoef_arrays = NULL;
  jpeg_saved_marker_ptr srcmarker_ptr;

  dstinfo->image_width = srcinfo->image_width;
  dstinfo->image_height = srcinfo->image_height;
  dstinfo->input_components = srcinfo->out_color_components;
  dstinfo->in_color_space = srcinfo->out_color_space;
  jpeg_set_defaults(dstinfo);
  dstinfo->optimize_coding = 0;
  dstinfo->num_components = srcinfo->num_components;
  dstinfo->jpeg_color_space = srcinfo->jpeg_color_space;
  if (srcinfo->progressive_mode) {
      jpeg_simple_progression(dstinfo);
  }
  for (ci = 0; ci < dstinfo->num_components; ci++) {
    dstinfo->comp_info[ci].component_id = srcinfo->comp_info[ci].component_id;
    dstinfo->comp_info[ci].h_samp_factor = srcinfo->comp_info[ci].h_samp_factor;
    dstinfo->comp_info[ci].v_samp_factor = srcinfo->comp_info[ci].v_samp_factor;
    dstinfo->comp_info[ci].quant_tbl_no = srcinfo->comp_info[ci].quant_tbl_no;
    dstinfo->comp_info[ci].ac_tbl_no = srcinfo->comp_info[ci].ac_tbl_no;
    dstinfo->comp_info[ci].dc_tbl_no = srcinfo->comp_info[ci].dc_tbl_no;
  }

  for (n = 0; n < NUM_QUANT_TBLS; n++) {
	  if (srcinfo->quant_tbl_ptrs[n] != NULL) {
		  JQUANT_TBL *quant_ptr;
		  if (dstinfo->quant_tbl_ptrs[n] == NULL)
			  dstinfo->quant_tbl_ptrs[n] = jpeg_alloc_quant_table((j_common_ptr) dstinfo);

		  quant_ptr = srcinfo->quant_tbl_ptrs[n];

		  for (i = 0; i < DCTSIZE; i++) 
			  for (j = 0; j < DCTSIZE; j++) {
				  int t = quant_ptr->quantval[i*DCTSIZE+j];
				  if (t<1 || t>65535)
					  fprintf(stderr, "Quantization table entries not in range 1..65535");
				  dstinfo->quant_tbl_ptrs[n]->quantval[i*DCTSIZE+j] = (UINT16) t;
			  }
	  }
	  else
		  dstinfo->quant_tbl_ptrs[n] = NULL;
  }

  srccoef_arrays = jpeg_read_coefficients(srcinfo);
  dstcoef_arrays = (jvirt_barray_ptr *) 
    (dstinfo->mem->alloc_small) ((j_common_ptr) dstinfo, JPOOL_IMAGE, sizeof(jvirt_barray_ptr) * dstinfo->num_components);
  for (ci = 0; ci < srcinfo->num_components; ci++) {
    jpeg_component_info *dstcompptr = dstinfo->comp_info + ci;
    jpeg_component_info *srccompptr = srcinfo->comp_info + ci;
    dstcompptr->height_in_blocks = srccompptr->height_in_blocks;
	dstcompptr->width_in_blocks = srccompptr->width_in_blocks;
  
    dstcoef_arrays[ci] = (dstinfo->mem->request_virt_barray) 
		((j_common_ptr) dstinfo, JPOOL_IMAGE, TRUE,
       (JDIMENSION) jround_up((long) dstcompptr->width_in_blocks,
                              (long) dstcompptr->h_samp_factor),
       (JDIMENSION) jround_up((long) dstcompptr->height_in_blocks,
                              (long) dstcompptr->v_samp_factor),
       (JDIMENSION) dstcompptr->v_samp_factor);
  }
  jpeg_write_coefficients(dstinfo,dstcoef_arrays);

  if (dstinfo->optimize_coding == FALSE) {
    for (n = 0; n < NUM_HUFF_TBLS; n++) {
      if (srcinfo->ac_huff_tbl_ptrs[n] != NULL) {
		  JHUFF_TBL *srchuff_ptr = NULL;
          if (dstinfo->ac_huff_tbl_ptrs[n] == NULL)
			  dstinfo->ac_huff_tbl_ptrs[n] = jpeg_alloc_huff_table((j_common_ptr) dstinfo);
	      srchuff_ptr = srcinfo->ac_huff_tbl_ptrs[n];
          for (i = 1; i <= 16; i++)
				dstinfo->ac_huff_tbl_ptrs[n]->bits[i] = srchuff_ptr->bits[i];
		  for (i = 0; i < 256; i++)
				dstinfo->ac_huff_tbl_ptrs[n]->huffval[i] = srchuff_ptr->huffval[i];
      }
	  else 
		  dstinfo->ac_huff_tbl_ptrs[n] = NULL;

      if (srcinfo->dc_huff_tbl_ptrs[n] != NULL) {
			JHUFF_TBL *srchuff_ptr = NULL;
			if (dstinfo->dc_huff_tbl_ptrs[n] == NULL)
			  dstinfo->dc_huff_tbl_ptrs[n] = jpeg_alloc_huff_table((j_common_ptr) dstinfo);
	        srchuff_ptr = srcinfo->dc_huff_tbl_ptrs[n];
			for (i = 1; i <= 16; i++)
				dstinfo->dc_huff_tbl_ptrs[n]->bits[i] = srchuff_ptr->bits[i];
			for (i = 0; i < 256; i++)
				dstinfo->dc_huff_tbl_ptrs[n]->huffval[i] = srchuff_ptr->huffval[i];
      }
	  else 
		  dstinfo->dc_huff_tbl_ptrs[n] = NULL;
	}
  }
  srcmarker_ptr = srcinfo->marker_list;
  while (srcmarker_ptr != NULL) {
    switch (srcmarker_ptr->marker) {
		case JPEG_COM:
			jpeg_write_marker(dstinfo, JPEG_COM, srcmarker_ptr->data, srcmarker_ptr->data_length);
			break;
		default:
			break;
    }
    srcmarker_ptr = srcmarker_ptr->next;
  }

  // copy DCT coef
  for (ci = 0; ci < dstinfo->num_components; ci++) {
    jpeg_component_info *dstcompptr = dstinfo->comp_info + ci;
	//jpeg_component_info *srccompptr = srcinfo->comp_info + ci;
	JDIMENSION blk_x, blk_y;
    for (blk_y = 0; blk_y < dstcompptr->height_in_blocks; blk_y++) {
      JBLOCKARRAY srcbuffer = (srcinfo->mem->access_virt_barray)
    	((j_common_ptr) srcinfo, srccoef_arrays[ci], blk_y, 1, FALSE);
      JBLOCKARRAY dstbuffer = (dstinfo->mem->access_virt_barray)
	  	((j_common_ptr) dstinfo, dstcoef_arrays[ci], blk_y, 1, TRUE);

	  for (blk_x = 0; blk_x < dstcompptr->width_in_blocks; blk_x++) {
        JCOEFPTR srcbufptr = srcbuffer[0][blk_x];
        JCOEFPTR dstbufptr = dstbuffer[0][blk_x];
		for (i = 0; i < DCTSIZE2; i++)
			dstbufptr[i] = srcbufptr[i];
	  }
	}
  }

  return dstcoef_arrays;
}

long countNonZeroOne(jvirt_barray_ptr *coef_arrays, struct jpeg_compress_struct *dstinfo) {
	long count = 0;
	int ci, i;

	for (ci = 0;ci < dstinfo->num_components;ci++) {
		JDIMENSION height = dstinfo->comp_info[ci].height_in_blocks;
		JDIMENSION width = dstinfo->comp_info[ci].width_in_blocks;
		JDIMENSION blk_x, blk_y;
		for (blk_y = 0; blk_y < height; blk_y++) {
			JBLOCKARRAY buffer = (dstinfo->mem->access_virt_barray) ((j_common_ptr) dstinfo, coef_arrays[ci], blk_y, 1, TRUE);
			for (blk_x = 0; blk_x < width; blk_x++) {
				JCOEFPTR bufptr = buffer[0][blk_x];
				for (i = 0; i < DCTSIZE2; i++) {
					short p = (short)(bufptr[i]);
					if ((p != 0) && (p != 1))
						count++;
				}
			}
		}
	}
	return count;
}

long countNonZeroOne_d(jvirt_barray_ptr *coef_arrays, struct jpeg_decompress_struct *srcinfo) {
	long count = 0;
	int ci, i;
	for (ci = 0;ci < srcinfo->num_components;ci++) {
		JDIMENSION height = srcinfo->comp_info[ci].height_in_blocks;
		JDIMENSION width = srcinfo->comp_info[ci].width_in_blocks;
		JDIMENSION blk_x, blk_y;
		for (blk_y = 0; blk_y < height; blk_y++) {
			JBLOCKARRAY buffer = (srcinfo->mem->access_virt_barray) ((j_common_ptr) srcinfo, coef_arrays[ci], blk_y, 1, TRUE);
			for (blk_x = 0; blk_x < width; blk_x++) {
				JCOEFPTR bufptr = buffer[0][blk_x];
				for (i = 0; i < DCTSIZE2; i++) {
					short p = (short)(bufptr[i]);
					if ((p != 0) && (p != 1))
						count++;
				}
			}
		}
	}
	return count;
}

long countNonZero(jvirt_barray_ptr *coef_arrays, struct jpeg_compress_struct *dstinfo) {
	long count = 0;
	int ci, i;
	for (ci = 0;ci < dstinfo->num_components;ci++) {
		JDIMENSION height = dstinfo->comp_info[ci].height_in_blocks;
		JDIMENSION width = dstinfo->comp_info[ci].width_in_blocks;
		JDIMENSION blk_y, blk_x;
		for (blk_y = 0; blk_y < height; blk_y++) {
			JBLOCKARRAY buffer = (dstinfo->mem->access_virt_barray) ((j_common_ptr) dstinfo, coef_arrays[ci], blk_y, 1, TRUE);
			for (blk_x = 0; blk_x < width; blk_x++) {
				JCOEFPTR bufptr = buffer[0][blk_x];
				// skip DC
				for (i = 1; i < DCTSIZE2; i++) {
					short p = (short)(bufptr[i]);
					if (p != 0)
						count++;
				}
			}
		}
	}
	return count;
}

long countNonZero_d(jvirt_barray_ptr *coef_arrays, struct jpeg_decompress_struct *srcinfo) {
	long count = 0;
	int ci, i;
	for (ci = 0;ci < srcinfo->num_components;ci++) {
		JDIMENSION height = srcinfo->comp_info[ci].height_in_blocks;
		JDIMENSION width = srcinfo->comp_info[ci].width_in_blocks;
		JDIMENSION blk_y, blk_x;
		for (blk_y = 0; blk_y < height; blk_y++) {
			JBLOCKARRAY buffer = (srcinfo->mem->access_virt_barray) ((j_common_ptr) srcinfo, coef_arrays[ci], blk_y, 1, TRUE);
			for (blk_x = 0; blk_x < width; blk_x++) {
				JCOEFPTR bufptr = buffer[0][blk_x];
				// skip DC
				for (i = 1; i < DCTSIZE2; i++) {
					short p = (short)(bufptr[i]);
					if (p != 0)
						count++;
				}
			}
		}
	}
	return count;
}

#endif